import asyncio

from pubsub.utils import ensure_bytes
from .base import BaseBroker


class RedisBroker(BaseBroker):
    def __init__(self, pool, *args, **kwargs):
        super().__init__(pool, *args, **kwargs)
        self.pool = pool
        self.connection = None
        self.background_tasks = {}
        self.channels_subscribed = []

    async def get_connection(self):
        if not self.connection:
            self.connection = await self.pool.acquire()
            self.logger.debug(f'get connection {self.connection} from pool')
        return self.connection

    async def start_publish(self):
        pass

    async def start_consume(self, channels_subscribed, callback):
        await self.get_connection()
        self.channels_subscribed = channels_subscribed
        for channel in channels_subscribed:
            await self.connection.execute_pubsub('subscribe', channel)
            ch = self.connection.pubsub_channels[channel]
            task = asyncio.create_task(self._subscribe_task(ch, callback))
            self.background_tasks[channel] = task
            self.logger.debug(f'background task {task} is created')
            self.logger.debug('subscribe channel %s', channel)

    async def _subscribe_task(self, ch, callback):
        while await ch.wait_message():
            decoded_body = await ch.get()
            self.logger.debug(f'[consume_message] {decoded_body}')
            if callback is None:
                self.logger.warning('[consume_message] callback is None')
                return
            try:
                await callback(decoded_body)
            except Exception as err:
                self.logger.exception(
                    '[consume_message] unexpected err: %s',
                    err,
                    exc_info=True,
                )

    async def stop(self):
        if self.channels_subscribed:
            for channel in self.channels_subscribed:
                self.logger.debug('unsubscribe channel %s', channel)
                await self.connection.execute_pubsub('unsubscribe', channel)
        for task in self.background_tasks.values():
            if task.done() or task.cancelled():
                continue
            task.cacel()
            self.logger.debug('background task {task} is canceled')
        self.background_tasks = {}
        if self.connection:
            self.pool.release(self.connection)

    async def publish(self, message, *channel_args, **channel_kwargs):
        channels = channel_args
        assert isinstance(channels, (tuple, list))
        message = ensure_bytes(message)
        for channel in channels:
            await self.pool.execute('PUBLISH', channel, message)
